#!/usr/bin/env python
import sys
import math
import numpy as np
import os, time, random
from multiprocessing import Pool

def process_images(name, task_cmd):
    print('Run task %s (%s)...' % (name, os.getpid()))
    start = time.time()
    print('[RUN CMD]: %s' % (task_cmd))
    os.system(task_cmd)
    end = time.time()
    print('Task %s runs %0.2f seconds.' % (name, (end - start)))

def get_all_task_infos(src_task_file_dir, origin_t1_dir, processed_t1_dir):
    file_list = []
    for file_name in os.listdir(src_task_file_dir):
        if "recon_list" in file_name and "t1.lock" in file_name:
            file_list.append(os.path.join(src_task_file_dir, file_name))

    subname_list = []
    for sub_file in file_list:
        fp = open(sub_file, "r")

        for line in fp.readlines():
            subj_name = line.replace("\n", "")
            subname_list.append(subj_name)

        fp.close()

    cmd_list = []
    for subj_name in subname_list:
        input_dir = os.path.join(origin_t1_dir, subj_name)
        iput_file = os.listdir(input_dir)[0]
        full_input_file = os.path.join(input_dir, iput_file)

        recon_cmd = "recon-all"
        # -s
        recon_cmd = recon_cmd + " -s " + subj_name
        # -sd
        recon_cmd = recon_cmd + " -sd " + processed_t1_dir
        # qchche
        recon_cmd = recon_cmd + " -qcache"
        cmd_list.append(recon_cmd)

    return cmd_list, len(cmd_list)

if __name__=='__main__':
    print('Parent process %s.' % os.getpid())
    process_num = 11

    project_dir = '/mri_projects/ASD_Analysis'

    src_task_file_dir = os.path.join(project_dir, 'Lists')
    origin_t1_dir = os.path.join(project_dir, 'Origin/T1/')
    processed_t1_dir = os.path.join(project_dir, 'DataProcessed/T1/')

    src_lock_file = os.path.join(src_task_file_dir, "t1.qchche.lock")
    if os.path.exists(src_lock_file):
        print("[Notice]: File List in `%s` Qcache Finished!" % (src_task_file_dir))
    else:
        if not os.path.exists(processed_t1_dir):
            os.makedirs(processed_t1_dir)

        task_list, task_size = get_all_task_infos(src_task_file_dir, origin_t1_dir, processed_t1_dir)

        p = Pool(process_num)
        for i in range(task_size):
            p.apply_async(process_images, args=(i,task_list[i],))
        print('Waiting for all subprocesses done...')
        p.close()
        p.join()
        print('All subprocesses done.')

        # print("########## Lock File ##########")
        # print("[Lock T1]: %s" % (src_task_file))
        # cp_cmd = "cp " + src_task_file + " " + src_lock_file
        # os.system(cp_cmd)
        # print("########## Done ##########")
